import os
import sys
import json
import argparse

# 导入必要的函数
parent_path = os.path.abspath(os.path.join(__file__, *(['..'] * 2)))
sys.path.insert(0, parent_path)
from _base_.base import train
from _base_.dataset_utils import DatasetConfigLoader

# 定义训练主函数
def main():
    parser = argparse.ArgumentParser(description='Train mask-rcnn detector')
    parser.add_argument('config_path', help='train config file data (json format)')
    parser.add_argument('output_dir', help='specify output directory')
    args = parser.parse_args()

    with open(args.config_path, 'rt', encoding='utf-8') as f:
        train_cfg = json.load(f)
    base_cfg_path = os.path.join( parent_path, 'configs/mask_rcnn/%s.py' % train_cfg['config_name'] )
    config_loader = DatasetConfigLoader()
    num_classes, user_cfg_dict = config_loader.load_json_data(train_cfg)

    # 我们还需要更改 head 中的 num_classes 以匹配数据集中的类别数
    user_cfg_dict['model'] = dict(
        roi_head=dict( 
            bbox_head=dict(num_classes=num_classes),
            mask_head=dict(num_classes=num_classes)
        ))
    
    train( base_cfg_path, user_cfg_dict, args.output_dir )

main()
